barebones LSP Server implementation by KRRT7 · Pull Request #309 · codeflash-ai/codeflash

KRRT7 · 2025-06-09T17:23:54Z

PR Type

Enhancement

Description

Integrate LSP server for function optimization
Define LSP methods for optimization workflow
Modularize & refactor FunctionOptimizer methods
Introduce benchmark and test discovery utilities

Changes walkthrough 📝

Relevant files

Enhancement

beta.py `Implement LSP optimization endpoints` codeflash/lsp/beta.py Add LSP feature handlers for optimization Implement getOptimizableFunctions and initialization Support test discovery, preparation, and execution Define performFunctionOptimization endpoint	+231/-0
server.py `Add custom LSP server protocol and initialization` codeflash/lsp/server.py Define custom LanguageServerProtocol with initialize Locate and load `pyproject.toml` for optimizer Implement CodeflashLanguageServer initializer	+56/-0
function_optimizer.py `Modularize FunctionOptimizer methods` codeflash/optimization/function_optimizer.py Extract `can_be_optimized` helper method Add modular test generation and baseline setup Refactor `optimize_function` to use new steps	+296/-160
optimizer.py `Enhance Optimizer with LSP integration` codeflash/optimization/optimizer.py Introduce `run_benchmarks` and `discover_tests` Add `prepare_module_for_optimization` helper Update `create_function_optimizer` signature Maintain current function context for LSP	+186/-153

Dependencies

pyproject.toml `Update dependencies and lint ignore` pyproject.toml Add direct dependency on `pygls` Update lint ignore rules	+3/-1

Additional files

__init__.py	[link]

Need help?
Type /help how to ... in the comments thread for any questions about PR-Agent usage.
Check out the documentation for more information.

…extension`) Here's an optimized version of your program. Main changes. - Replace slices and length-based for-loops with more direct, idiomatic Python control flow (they did nothing). - Combine checks for better early returns. - Reduce unnecessary list accesses. - Process decorator list without unnecessary nesting. - No changes to function signatures or output. All your in-code comments are preserved (there were none). **Key improvements:** - Direct mapping of `parents[0].type`; no useless range-loop. - Early exit for empty `parents`. - Avoid redundant list traversal. - Set literal for type checking. This function has the same return signature and logic, but runs faster and is easier to reason about.

codeflash-ai · 2025-06-16T21:30:01Z

⚡️ Codeflash found optimizations for this PR

📄 39% (0.39x) speedup for `function_kind` in `codeflash/code_utils/static_analysis.py`

⏱️ Runtime : 1.77 milliseconds → 1.27 milliseconds (best of 158 runs)

I created a new dependent PR with the suggested changes. Please review:

⚡️ Speed up function function_kind by 39% in PR #309 (updated-vsc-extension) #330

If you approve, it will be merged into this PR (branch updated-vsc-extension).

github-actions · 2025-06-25T23:19:24Z

PR Reviewer Guide 🔍

(Review updated until commit `89ceca8`)

Here are some key observations to aid the review process:

⏱️ Estimated effort to review: 4 🔵🔵🔵🔵⚪
🧪 No relevant tests
🔒 No security concerns identified
⚡ Recommended focus areas for review Missing None-handling LSP handlers like `prepareOptimization` and `performFunctionOptimization` unpack the result of `prepare_module_for_optimization` without checking for `None`, which can lead to runtime errors if parsing fails. @server.feature("prepareOptimization") def prepare_optimization(server: CodeflashLanguageServer, params: FunctionOptimizationParams) -> dict[str, str]: current_function = server.optimizer.current_function_being_optimized module_prep_result = server.optimizer.prepare_module_for_optimization(current_function.file_path) validated_original_code, original_module_ast = module_prep_result function_optimizer = server.optimizer.create_function_optimizer( current_function, function_to_optimize_source_code=validated_original_code[current_function.file_path].source_code, original_module_ast=original_module_ast, original_module_path=current_function.file_path, ) Unimplemented method The newly declared `submit_test_generation_tasks` method is referenced in test generation but its implementation is not shown, likely causing failures at runtime. def submit_test_generation_tasks( self, Stale state The attribute `generate_and_instrument_tests_results` is set but never cleared or used elsewhere, which may retain outdated data between optimization runs. self.generate_and_instrument_tests_results: ( tuple[GeneratedTestsList, dict[str, set[FunctionCalledInTest]], OptimizationSet] \| None ) = None

github-actions · 2025-06-25T23:20:20Z

PR Code Suggestions ✨

Latest suggestions up to 89ceca8
Explore these optional code suggestions:

Category	Suggestion	Impact
Possible issue	Define missing context before returning The variables `code_context` and `original_helper_code` are not defined, causing a NameError. You should call `prepare_module_for_optimization` to obtain the validated code and AST, then construct a `CodeOptimizationContext` and extract the helper code mapping before returning Success. codeflash/optimization/function_optimizer.py [150-157] def can_be_optimized(self) -> Result[tuple[bool, CodeOptimizationContext, dict[Path, str]], str]: + prep = self.prepare_module_for_optimization(self.function_to_optimize.file_path) + if prep is None: + return Failure("Module preparation failed") + validated_code, module_ast = prep + code_context = CodeOptimizationContext(validated_code, module_ast) + original_helper_code = {path: vc.source_code for path, vc in validated_code.items()} should_run_experiment = self.experiment_id is not None logger.debug(f"Function Trace ID: {self.function_trace_id}") ph("cli-optimize-function-start", {"function_trace_id": self.function_trace_id}) self.cleanup_leftover_test_return_values() - if random.random() > REPEAT_OPTIMIZATION_PROBABILITY and was_function_previously_optimized( self.function_to_optimize, code_context ): return Failure("Function optimization previously attempted, skipping.") - return Success((should_run_experiment, code_context, original_helper_code)) Suggestion importance[1-10]: 9 __ Why: The new `can_be_optimized` function references undefined `code_context` and `original_helper_code`, causing a NameError; fixing this is critical for the optimization pipeline.	High
	Handle failed module preparation `prepare_module_for_optimization` can return `None` on syntax errors, leading to a TypeError when unpacking. Add a check for `None` and return an LSP error response if module preparation failed. codeflash/lsp/beta.py [76-77] -@server.feature("prepareOptimization") -def prepare_optimization(server: CodeflashLanguageServer, params: FunctionOptimizationParams) -> dict[str, str]: - current_function = server.optimizer.current_function_being_optimized +module_prep_result = server.optimizer.prepare_module_for_optimization(current_function.file_path) +if module_prep_result is None: + return {"functionName": params.functionName, "status": "error", "message": "Module preparation failed"} +validated_original_code, original_module_ast = module_prep_result - module_prep_result = server.optimizer.prepare_module_for_optimization(current_function.file_path) - validated_original_code, original_module_ast = module_prep_result - Suggestion importance[1-10]: 7 __ Why: Adding a check for `None` prevents a TypeError when unpacking a failed `prepare_module_for_optimization`, improving robustness.	Medium
	Guard uninitialized current function If `current_function_being_optimized` is `None`, accessing `file_path` will raise an AttributeError. Add a guard to return an error when no function is initialized before discovery. codeflash/lsp/beta.py [61-65] -@server.feature("discoverFunctionTests") -def discover_function_tests(server: CodeflashLanguageServer, params: FunctionOptimizationParams) -> dict[str, str]: - current_function = server.optimizer.current_function_being_optimized +if not server.optimizer.current_function_being_optimized: + return {"functionName": params.functionName, "status": "error", "message": "No function initialized"} +current_function = server.optimizer.current_function_being_optimized +optimizable_funcs = {current_function.file_path: [current_function]} - optimizable_funcs = {current_function.file_path: [current_function]} - Suggestion importance[1-10]: 6 __ Why: Without the guard, `current_function` may be `None` and cause an AttributeError; adding this check prevents a potential runtime crash.	Low

Previous suggestions

Suggestions up to commit 3006d62

Category	Suggestion	Impact
Possible issue	Select function by name Instead of arbitrarily popping an item, filter the `optimizable_funcs` mapping by the requested function name. This ensures you select the correct function rather than an arbitrary one. Also handle the case where the function isn't found. codeflash/lsp/beta.py [56] -fto = optimizable_funcs.popitem()[1][0] +fto = next( + ( + func + for funcs in optimizable_funcs.values() + for func in funcs + if func.qualified_name == params.functionName + ), + None, +) +if not fto: + return {"functionName": params.functionName, "status": "not found", "info": None} Suggestion importance[1-10]: 8 __ Why: The code currently uses `popitem` arbitrarily which can select the wrong function; filtering by `params.functionName` ensures the correct function is chosen.	Medium
Possible issue	Prevent divide-by-zero error Guard against division by zero when computing speedup to avoid runtime errors. Check that `best_optimization.runtime` is non-zero before dividing. codeflash/lsp/beta.py [223] -"extra": f"Speedup: {original_code_baseline.runtime / best_optimization.runtime:.2f}x faster" +if best_optimization.runtime: + speedup = original_code_baseline.runtime / best_optimization.runtime + extra_msg = f"Speedup: {speedup:.2f}x faster" +else: + extra_msg = "Speedup: N/A" +... +"extra": extra_msg Suggestion importance[1-10]: 7 __ Why: Adding a guard against `best_optimization.runtime` being zero prevents potential runtime exceptions when calculating the speedup.	Medium
General	Return actual discovered tests Rather than returning the count as a string, return the actual number or the list of discovered tests to the client. This provides richer context and keeps types correct. codeflash/lsp/beta.py [69] -return {"functionName": params.functionName, "status": "success", "generated_tests": str(num_discovered_tests)} +discovered_tests = function_to_tests.get(current_function.file_path, []) +return { + "functionName": params.functionName, + "status": "success", + "discovered_tests": discovered_tests +} Suggestion importance[1-10]: 6 __ Why: Returning the actual list of discovered tests (or proper numeric type) provides more useful context and corrects the type mismatch.	Low
General	Consistent response schema Align the error response field names with the success case by using `"info"` instead of `"args"`. This keeps your response schema consistent. codeflash/lsp/beta.py [54-55] if not optimizable_funcs: - return {"functionName": params.functionName, "status": "not found", "args": None} + return {"functionName": params.functionName, "status": "not found", "info": None} Suggestion importance[1-10]: 5 __ Why: Changing the error response key from `args` to `info` matches the success case and keeps the response payload consistent.	Low

github-actions · 2025-06-26T00:26:50Z

Persistent review updated to latest commit 89ceca8

Saga4

LGTM

KRRT7 changed the title ~~Updated vsc extension~~ VSC Extension WIP Jun 9, 2025

KRRT7 force-pushed the updated-vsc-extension branch 2 times, most recently from 020b000 to 3c8d0e5 Compare June 11, 2025 20:28

codeflash-ai Bot mentioned this pull request Jun 16, 2025

⚡️ Speed up function function_kind by 39% in PR #309 (updated-vsc-extension) #330

Closed

KRRT7 closed this Jun 25, 2025

KRRT7 force-pushed the updated-vsc-extension branch from f2f394d to cef8503 Compare June 25, 2025 23:16

codeflash-ai Bot mentioned this pull request Jun 25, 2025

⚡️ Speed up function function_kind by 93% #323

Closed

add lSP module

3006d62

KRRT7 reopened this Jun 25, 2025

github-actions Bot added the Review effort 4/5 label Jun 25, 2025

KRRT7 and others added 11 commits June 25, 2025 16:23

extract benchmark runs

1311198

extract get_functions_to_optimize

30e51b8

extract create_function_optimizer

2ff091b

extract module prep

e48e13f

extract test discovery

e91838c

add pygls as a direct dependency, not transitive

f52ada5

implement can be optimized

ebe4f25

modularize optimize_function

0bc582f

fix regression

95fb1f0

regression fix

fa07342

Merge branch 'main' into updated-vsc-extension

89ceca8

KRRT7 changed the title ~~VSC Extension WIP~~ barebones LSP Server implementation Jun 26, 2025

KRRT7 marked this pull request as ready for review June 26, 2025 00:25

KRRT7 requested review from aseembits93 and misrasaurabh1 June 26, 2025 00:31

Merge branch 'main' into updated-vsc-extension

d0e6284

Saga4 approved these changes Jun 26, 2025

View reviewed changes

Merge branch 'main' into updated-vsc-extension

6347a8a

KRRT7 enabled auto-merge June 26, 2025 18:14

KRRT7 merged commit 4cf8ea5 into main Jun 26, 2025
16 checks passed

KRRT7 deleted the updated-vsc-extension branch June 26, 2025 18:18

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

barebones LSP Server implementation#309

barebones LSP Server implementation#309
KRRT7 merged 14 commits into
mainfrom
updated-vsc-extension

KRRT7 commented Jun 9, 2025 •

edited by github-actions Bot

Loading

Uh oh!

codeflash-ai Bot commented Jun 16, 2025

⚡️ Speed up function `function_kind` by 39% in PR #309 (`updated-vsc-extension`) #330

Uh oh!

github-actions Bot commented Jun 25, 2025 •

edited

Loading

Uh oh!

github-actions Bot commented Jun 25, 2025 •

edited

Loading

Uh oh!

github-actions Bot commented Jun 26, 2025

Uh oh!

Saga4 left a comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

KRRT7 commented Jun 9, 2025 • edited by github-actions Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

PR Type

Description

Changes walkthrough 📝

Uh oh!

codeflash-ai Bot commented Jun 16, 2025

⚡️ Codeflash found optimizations for this PR

📄 39% (0.39x) speedup for function_kind in codeflash/code_utils/static_analysis.py

I created a new dependent PR with the suggested changes. Please review:

⚡️ Speed up function function_kind by 39% in PR #309 (updated-vsc-extension) #330

Uh oh!

github-actions Bot commented Jun 25, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

PR Reviewer Guide 🔍

(Review updated until commit 89ceca8)

Uh oh!

github-actions Bot commented Jun 25, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

PR Code Suggestions ✨

Previous suggestions

Uh oh!

github-actions Bot commented Jun 26, 2025

Uh oh!

Saga4 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

KRRT7 commented Jun 9, 2025 •

edited by github-actions Bot

Loading

📄 39% (0.39x) speedup for `function_kind` in `codeflash/code_utils/static_analysis.py`

⚡️ Speed up function `function_kind` by 39% in PR #309 (`updated-vsc-extension`) #330

github-actions Bot commented Jun 25, 2025 •

edited

Loading

(Review updated until commit `89ceca8`)

github-actions Bot commented Jun 25, 2025 •

edited

Loading